package com.ls.controller;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.thread.ExecutorBuilder;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.jgit.api.Git;
import org.eclipse.jgit.api.errors.GitAPIException;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.FileSystemResource;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @BelongsProject: knowledge
 * @BelongsPackage: com.ls.controller
 * @Author: lishun
 * @CreateTime: 2025-10-27  23:17
 * @Description: 聊天测试
 * @Version: 1.0
 */
@RestController
@RequestMapping("chatTest")
@Slf4j
@CrossOrigin(origins = "*")
public class ChatController {

    private final ChatClient chatClient;


    private final VectorStore vectorStore;

    private final MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder().build();

    private final Map<String, SseEmitter> sseClientMap = new ConcurrentHashMap<>();


    // 使用线程池（推荐）
    private final ExecutorService executor = ExecutorBuilder.create()
            .setCorePoolSize(20)
            .setMaxPoolSize(20)
            .setWorkQueue(new LinkedBlockingQueue<>(10000))
            .build();


    public ChatController(ChatClient.Builder builder, EmbeddingModel embeddingModel) {
        this.vectorStore = SimpleVectorStore.builder(embeddingModel).build();
        this.chatClient = builder
                .defaultAdvisors(QuestionAnswerAdvisor.builder(vectorStore).build())
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build(), new SimpleLoggerAdvisor())
                .build();

    }

    @GetMapping("chat")
    public String chat(String msg) {
        String content = chatClient.prompt(msg).call().content();
        return content;
    }


    @GetMapping(value = "chatStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chatStream(String msg) {
        Flux<String> content = chatClient.prompt(msg).stream().content();
        return content;
    }

    @GetMapping(value = "chatMemoryTest", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chatMemory(String conversationId, String userMsg) {
        Flux<String> content = chatClient
                .prompt()
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
                .advisors(MessageChatMemoryAdvisor.builder(chatMemory).build(), new SimpleLoggerAdvisor())
                .user(userMsg)
                .stream().content();
        return content;
    }

    @GetMapping("initGitKnowledge")
    public void initGitKnowledge(String gitUrl) throws GitAPIException {
        File localPath = new File("git/" + UUID.randomUUID().toString()); // Windows 示例
        Git.cloneRepository()
                .setURI(gitUrl)
                .setDirectory(localPath)
                .setBranch("master")
                .call();
        FileUtil.walkFiles(localPath, file -> {
            executor.submit(() -> {
                TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(new FileSystemResource(file));
                List<Document> documents = tikaDocumentReader.get();
                String relativePath = localPath.toURI().relativize(file.toURI()).getPath();
                for (Document doc : documents) {
                    doc.getMetadata().put("gitUrl", gitUrl);
                    doc.getMetadata().put("sourceFile", relativePath);
                    doc.getMetadata().put("fileName", file.getName());
                }

                // 切割文本

                TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
                List<Document> documentList = tokenTextSplitter.apply(documents);

                if (CollectionUtil.isNotEmpty(documentList)) {
                    vectorStore.add(documentList);
                }
            });
        });


    }


    @GetMapping(value = "chatStreamRag", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public String chatStreamRag(@RequestParam("content") String content, @RequestParam("url") String url) {
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.90)
                        .vectorStore(vectorStore)
                        .build())
                .build();


        return this.chatClient.prompt()
                .user(content)
                .advisors(retrievalAugmentationAdvisor)
                .advisors(a -> a.param("gitUrl", url))
                .call().content();
    }


    @GetMapping(value = "/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter sseTest(String conversationId,String msg) {
        SseEmitter sseEmitter = new SseEmitter(3000000L);
        sseEmitter.onCompletion(() -> {
            log.info("连接关闭");
        });
        sseEmitter.onError((e) -> {
            log.info("连接错误", e);
        });
        sseEmitter.onTimeout(() -> {
            log.info("连接超时");
        });

        executor.submit(() -> {
            log.info("开始生成");
            Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                    .documentRetriever(VectorStoreDocumentRetriever.builder()
                            .similarityThreshold(0.8)
                            .topK(20)
                            .vectorStore(vectorStore)
                            .build())
                    .build();
            Flux<String> flux = this.chatClient.prompt()
                    .user(msg)
                    .system("你是一个AI智能助手，使用中文回答用户的问题")
                    .advisors(retrievalAugmentationAdvisor)
                    .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
                    .stream()
                    .content();
            flux.subscribe((s) -> {
                log.info("ai生成：{}", s);
                if (sseEmitter != null) {
                    try {
                        sseEmitter.send(s);
                    } catch (IOException e) {
                        log.error("生成错误", e);
                    }
                }
            }, (e) -> {
                log.error("生成错误", e);
            }, () -> {
                log.info("生成完成");
            });
        });
        return sseEmitter;
    }


    @GetMapping(value = "/test")
    public String test(String conversationId,String msg) {
        return "success";
    }


}
